#coding=utf-8
'''

    机器人数据类封装

'''

from basic_math import geolib,filterlib,forcelib
import time
import numpy as np
import copy

class Terminal(geolib.GeoEntity):
    '''
    机器人各个关键部位（例如头部，脚部等）基本描述
    其中，基类中的几何学参数(x,y,z,roll,pitch,yaw)代表
    '''
    #   静态变量（宏定义）
    def __init__(self,x,y,z,roll=0,pitch=0,yaw=0,datQueueLen=50,name="body_part", size=[0.12,0.18,0.09],color=(0.7,0.7,0.7,1)):
        '''
        初始化基类
        :param datQueueLen: 缓存数据长度
        '''
        #   时间戳
        super().__init__(x,y,z,roll,pitch,yaw)
        self.stamp = time.time()
        #   描述
        self.name=name
        self.size=size
        self.color=color
        #   滤波器
        self.filterLevel=2      #默认卡尔曼滤波
        self.__filterRotation=filterlib.Kalman_Filter3D()
        #   空间描述数据
        self.__accelRaw = np.array([0.0, 0.0, 0.0])  # 加速度数据（原始）
        self.__gyroRaw = np.array([0.0, 0.0, 0.0])   # 角速度数据（原始）
        self.__postureRaw = np.array([0.0, 0.0, 0.0])   #姿态数据（原始）
        self.__Raw=np.array([0.0, 0.0, 0.0])    #位置数据
        self.__G = 0

        #   缓存和校准数据
        self.__cacheLen=datQueueLen
        self.__postureCalib=np.array([0.0, 0.0, 0.0])   #姿态校准数据
        self.__positionCalib=np.array([0.0, 0.0, 0.0])  #位置校准数据
        self.__GCalib=1.0                               #重力加速度校准值

        #   数据队列(观测)
        self.__postureQueue=[]
        self.__positionQueue=[]
        self.__accelQueue=[]
        self.__gyroQueue=[]
        self.__GRawQueue=[]
        self.rotationQueue=[]
        self.ZMPQueue=[]

        for i in range(datQueueLen):
            self.__postureQueue.append([0.0,0.0,0.0])
            self.__positionQueue.append([0.0,0.0,0.0])
            self.__accelQueue.append([0.0,0.0,0.0])
            self.__gyroQueue.append([0.0,0.0,0.0])
            self.__GRawQueue.append(0.0)
            self.rotationQueue.append(0.0)
            self.ZMPQueue.append(geolib.GeoEntity())

        #   观测数据(经过数据过滤后)
        self.ZMP = geolib.GeoEntity()
        self.positionMonitored=np.zeros(3,float)
        self.rotationMonitored=np.zeros(3,float)
        #   模型计算数据
        self.positionKinematic = np.zeros(3, float)
        self.rotationKinematic = np.zeros(3, float)
        #   质心（等效）

    def CalibrateSensor(self):
        '''
        校准传感器
        :param flag:
        :return:
        '''
        #   校准姿态和位置
        self.__postureCalib=np.mean(np.array(self.__postureQueue),axis=0)
        self.__positionCalib=np.mean(np.array(self.__positionQueue),axis=0)
        #   校准重力加速度
        self.__GCalib=9.8/np.average(self.__GRawQueue)

    def UpdateFBData(self,jsonData):
        '''
        从json数据中更新基本几何学特征
        :return:
        '''
        now=time.time()
        defaultAccelRatio = 20.0 / 32768.0
        defaultGyroRatio = (250.0 * 3.14159265 / 180.0) / 32768.0

        self.__accelRaw[0] = jsonData["accel"]["x"]*defaultAccelRatio
        self.__accelRaw[1] = jsonData["accel"]["y"]*defaultAccelRatio
        self.__accelRaw[2] = jsonData["accel"]["z"]*defaultAccelRatio

        #   角速度数据
        self.__gyroRaw[0] = -jsonData["gyro"]["x"] * defaultGyroRatio
        self.__gyroRaw[1] = -jsonData["gyro"]["y"]*defaultGyroRatio
        self.__gyroRaw[2] = jsonData["gyro"]["z"]*defaultGyroRatio

        #   姿态队列获取(无滤波)
        self.__postureRaw = geolib.calculate_tilt_angles(self.__accelRaw[0], self.__accelRaw[1],self.__accelRaw[2])
        self.__postureQueue.append(self.__postureRaw)
        self.__postureQueue.pop(0)
        #   加速度队列获取
        self.__accelQueue.append(self.__accelRaw);self.__accelQueue.pop(0)
        #   旋转队列获取
        self.__gyroQueue.append(self.__gyroRaw);self.__gyroQueue.pop(0)
        #   加速度队列获取
        gravity=np.linalg.norm(self.__accelRaw)
        self.__GRawQueue.append(gravity)
        self.__GRawQueue.pop(0)
        self.__G=gravity*self.__GCalib
        #   计算姿态（滤波）
        if self.filterLevel==0: self.rotationMonitored=self.__postureRaw - self.__postureCalib
        if self.filterLevel>=1:
            #   中值滤波
            self.rotationMonitored = filterlib.median_filter_1d(self.__postureQueue, 5)- self.__postureCalib
        if self.filterLevel>=2:
            #   卡尔曼滤波
            self.__filterRotation.Predict(velocity=self.__gyroRaw*(now-self.stamp))
            self.rotationMonitored=self.__filterRotation.Update(observed_Pos=self.rotationMonitored)
        self.rotationQueue.append(self.rotationMonitored)
        self.rotationQueue.pop(0)
        self.stamp=now

    def GetAccelRaw(self):
        '''
        获取原始加速度和
        :return:
        '''
        return self.__accelRaw
    def GetGyroRaw(self):
        '''
        获取原始旋转速度
        :return:
        '''
        return self.__gyroRaw

    def GetRotationError(self):
        '''
        观测值误差-理论误差
        :return:
        '''
        return self.rotationMonitored[0]-self.roll,self.rotationMonitored[1]-self.pitch,self.rotationMonitored[2]-self.yaw
    def CopyTerminal(self):
        '''
        生成一个几何属性完全一致的实体
        :return:
        '''
        return Terminal(self.x,self.y,self.z,self.roll,self.pitch,self.yaw,name='copy_'+self.name)
class BaseLink(Terminal):
    '''
    基座类
    '''

    def __init__(self,x,y,z,roll=0,pitch=0,yaw=0,datQueueLen=30,name="",size=[0.09,0.16,0.06],color=(0.7,0.7,0.7,1)):
        #   初始化几何学特性
        super().__init__(x,y,z,roll,pitch,yaw,datQueueLen,name,size,color)
        #   电量特征
        self.__batteryQueue=[]
        self.__powerQueue=[]
        self.__tempQueue=[]

        for i in range(datQueueLen):
            self.__batteryQueue.append(0)
            self.__powerQueue.append(0)
            self.__tempQueue.append(0)

    def UpdateFBData(self, jsonData, filterKernel=5):
        '''
        从json数据包中更新数据
        :param jsonData:
        :return:
        '''
        #   更新几何学数据
        super().UpdateFBData(jsonData)
        adc=jsonData["volt"]
        self.__batteryQueue.append(adc)
        self.__batteryQueue.pop(0)
        amp = jsonData["amp"]
        self.__powerQueue.append(amp)
        self.__powerQueue.pop(0)
        self.__tempQueue.append(jsonData['temp'])
        self.__tempQueue.pop(0)
    def ADC2Volt(self,val):
        '''
        获得真实电压值
        :param val:
        :return:
        '''
        #print(val)
        return 7.2+(val-1100)/400

    def GetBattery(self):
        '''
        获得电池电量以及电压
        :return:
        '''
        volt=np.average(np.array(self.__batteryQueue))
        percentage=100*(volt-7.2)/(8.4-7.2)
        amp=self.__powerQueue[-1]
        return percentage,amp

    def GetTemperature(self):
        temp=np.average(np.array(self.__tempQueue))
        return temp

class FootTerminal(Terminal):
    '''
    脚部末端数据定义
    基于几何学类
    含有压力传感器组，识别
    '''

    def __init__(self,x,y,z,roll=0,pitch=0,yaw=0,datQueueLen=30,name="",size=[0.18,0.09,0.06],color=(0.7,0.7,0.7,1),sensorPositions=[[0.04,0.02375],[0.04,-0.02375],[-0.04,0.02375],[-0.04,-0.02375]],free_threshold=2.0):
        #   初始化几何学特性
        super().__init__(x,y,z,roll,pitch,yaw,datQueueLen,name,size,color)
        #   压力传感器位置信息
        self.sensorPositions=sensorPositions
        self.sensorNums=len(sensorPositions)
        #   压力传感器数据
        self.__pressureRaw=np.zeros(4,float)
        self.__freeThres=free_threshold
        #   队列构建
        self.__pressureQueue=[]     #四个压力队列
        for i in range(datQueueLen):
            self.__pressureQueue.append(np.zeros(4))
        #   开放数据
        self.pressure=np.zeros(4)
        self.pressureTotal=0.0
        self.isSupport=False

    def UpdateFBData(self,jsonData,filterKernel=3):
        '''
        从json数据包中更新数据
        :param jsonData:
        :return:
        '''
        #   更新几何学数据
        super().UpdateFBData(jsonData)
        #   更新压力数据
        self.__pressureRaw[0]=jsonData["pressure"]["ADC0"]
        self.__pressureRaw[1] = jsonData["pressure"]["ADC1"]
        self.__pressureRaw[2] = jsonData["pressure"]["ADC2"]
        self.__pressureRaw[3] = jsonData["pressure"]["ADC3"]
        self.__pressureQueue.append(self.__pressureRaw)
        self.__pressureQueue.pop(0)
        p=filterlib.median_filter_1d(self.__pressureQueue,filterKernel)
        #   计算阻抗数值
        V_matrix = p * 3.3 / 4096
        R = V_matrix * 51.0 / (3.3 - V_matrix)
        self.pressure = 45.48 * np.power(R, -1)
        self.pressureTotal = np.sum(self.pressure)
        if self.pressureTotal<self.__freeThres: self.isSupport=False
        else:self.isSupport=True
        self.UpdateZMP()
    def UpdateZMP(self):
        '''
        更新ZMP点
        :return:
        '''
        self.ZMP.x,self.ZMP.y=forcelib.Sensor2ZMP(self.sensorPositions,self.pressure)
        self.ZMPQueue.append(copy.copy(self.ZMP))
        self.ZMPQueue.pop(0)

class AppleChain:
    '''
        各个Terminal之间的骨骼链条
        解算坐标系：Global
        这个是基于AppleZhang的固定模型进行解算的，所以我们把它叫做AppleChain,这个Chain中各个肢体的长度可以调整，但是模型是固定的
        每一个从base_link到足部Terminal之间，都由一条AppleChain进行连接
        ______
         头部
        ______
        |    |  （竖）
        chain[0]
        |    |  （横）
        chain[1]
        |    |  （纵）
        chain[2]
        |    |  （纵）
        chain[3]
        |    |  （纵）
        chain[4]
        |    |  （横）
        chain[5]
       -脚- -脚-
    '''
    def __init__(self,baseEntity:Terminal,terminalEntity:FootTerminal,baseTrans=np.zeros(3),boneLength=[0.043,0.1205,0.1205,0.033,0.0304]):
        '''
        :param baseEntity: 基座对象
        :param terminalEntity: 末端对象
        初始化Chain
        :param baseTrans:   从基座（两脚之间）到第一个关节的固定位置
        :param boneLength:  所有骨骼点的长度
        '''
        self.base=baseEntity
        self.terminal=terminalEntity
        self.baseTrans=baseTrans
        self.boneLength=boneLength
        self.boneNums=len(boneLength)
        #   关节数值描述
        self.joints=np.zeros(6)
        #   构建初始链（逆解姿态）
        self.chain=[]    #   逆解关节链
        #   链属性
        self.maxLength=0.0

        for i in range(6):
            self.chain.append(geolib.GeoLinks(name="[{}]".format(i),length=boneLength[i]))
            self.maxLength+=boneLength[i]

    def Updatechain(self,thresholds=[[-0.5,0.5],[-1.5,1.5],[-1.57,1.57],[-2.09,2.09],[-2.09,2.09],[-1.57,1.57]],autoAdjust="base"):
        '''
        在base和terminal之间构建链条（如果有的话）

        :return:
        '''
        #   解算第5个Chain的姿态
        self.chain[5].start = copy.copy(self.terminal)
        self.chain[5].UpdateEnd()



        #   第0个joint的旋转角度
        joints=np.array([0,0,0,0,0,0],dtype=float)
        joints[0] = self.terminal.yaw - self.base.yaw
        #   第0个Chain：起点为base_link，末端根据姿态计算
        self.chain[0].end=self.base.BoneTrans(x=self.baseTrans[0],y=self.baseTrans[1],l=self.baseTrans[2])
        self.chain[0].start.Parallel(self.base)
        self.chain[0].UpdateStart(yaw=joints[0])
        '''
            这里暂时不考虑yaw的成分
        '''
        x,y,z=self.chain[5].end.DistancebyAxis(self.chain[0].start) #得到两个点的差值
        plateZ=np.sqrt(y*y+z*z)
        plateY=y
        plateTheta=np.arctan(plateY/plateZ)
        joints[1]=plateTheta-self.chain[0].start.roll
        self.chain[1].end=copy.copy(self.chain[0].start)
        self.chain[1].UpdateStart(roll=joints[1])
        '''
            直接生成第5关节数值
        '''
        joints[5]=-plateTheta+self.terminal.roll
        self.chain[4].start=copy.copy(self.chain[5].end)
        self.chain[4].UpdateEnd(roll=-joints[5])
        '''
            接下来解算[2][3][4]关节数值
        '''
        x, y, z = self.chain[4].end.DistancebyAxis(self.chain[1].start)  # 得到两个点的差值
        distance=self.chain[4].end.Distance(self.chain[1].start)    #计算两个P点距离
        plateZ=np.sqrt(distance*distance-x*x)   #等效z轴
        plateX=x
        l=self.chain[2].lengh + self.chain[3].lengh
        if distance>l:
            '''
                无解
            '''
            if autoAdjust=='base':
                zMove=-z-np.sqrt(l*l-x*x-y*y)
                #   减去一个很小的数值以避免依然无解
                self.base.z-=(zMove+0.000001)
                #   使用递归实现更新坐标后继续逆解
                return self.Updatechain(autoAdjust='base')
            elif autoAdjust=='terminal':
                zMove = -z - np.sqrt(l * l - x * x - y * y)
                #   减去一个很小的数值以避免依然无解
                self.terminal.z += zMove + 0.000001
                #   使用递归实现更新坐标后继续逆解
                return self.Updatechain(autoAdjust='terminal')
            else:
                return False
        plateTheta=np.arctan(plateX/plateZ)
        theta1,theta2,theta3=geolib.calculate_triangle_angles(self.chain[3].lengh,distance,self.chain[2].lengh)
        #   更新J2以及chain[2]
        joints[2]=(-plateTheta-theta1)-self.chain[1].start.pitch

        self.chain[2].end=copy.copy(self.chain[1].start)
        self.chain[2].UpdateStart(pitch=joints[2])
        #   更新J3以及Chain3
        joints[3]=np.pi-theta2
        self.chain[3].end=copy.copy(self.chain[2].start)
        self.chain[3].UpdateStart(pitch=joints[3])
        joints[4]=plateTheta-theta3-self.terminal.pitch
        for i in range(6):
            if joints[i]>thresholds[i][1] or joints[i]<thresholds[i][0]:
                return False
        self.joints=joints
        self.chain=self.chain
        return True

    def Updatechain(self,autoAdjust='base'):
        '''
        从关节更新
        :param const:
        :return:
        '''
        if autoAdjust=='terminal':
            '''
            自动调整脚部坐标
            '''
            self.chain[0].end = self.base.BoneTrans(x=self.baseTrans[0], y=self.baseTrans[1], l=self.baseTrans[2])
            self.chain[0].UpdateStart(yaw=self.joints[0])
            self.chain[0].start.Parallel(self.base)
            self.chain[1].end = copy.copy(self.chain[0].start)
            self.chain[1].UpdateStart(roll=self.joints[1])
            self.chain[2].end = copy.copy(self.chain[1].start)
            self.chain[2].UpdateStart(pitch=self.joints[2])
            self.chain[3].end = copy.copy(self.chain[2].start)
            self.chain[3].UpdateStart(pitch=self.joints[3])
            self.chain[4].end = copy.copy(self.chain[3].start)
            self.chain[4].UpdateStart(pitch=self.joints[4])
            self.chain[5].end = copy.copy(self.chain[4].start)
            self.chain[5].UpdateStart(roll=self.joints[5])
            #   解算Terminal
            self.terminal.SetGeoData(self.chain[5].start.VectorPose(),self.chain[5].start.VectorRot())
            self.baseFb.SetGeoData(self.base.VectorPose(),rot=False)
            self.terminalFb.SetGeoData(self.terminal.VectorPose(),rot=False)
        else:
            '''
            自动调整基座
            '''
            self.chain[5].start=copy.copy(self.terminal)
            self.chain[5].UpdateEnd()
            self.chain[4].start=copy.copy(self.chain[5].end)
            self.chain[4].UpdateEnd(roll=-self.joints[5])
            self.chain[3].start = copy.copy(self.chain[4].end)
            self.chain[3].UpdateEnd(pitch=-self.joints[4])
            self.chain[2].start = copy.copy(self.chain[3].end)
            self.chain[2].UpdateEnd(pitch=-self.joints[3])
            self.chain[1].start = copy.copy(self.chain[2].end)
            self.chain[1].UpdateEnd(pitch=-self.joints[2])
            self.chain[0].start = copy.copy(self.chain[1].end)
            self.chain[0].UpdateEnd(roll=-self.joints[1])
            #   解算base_link
            basePos=np.array(self.chain[0].end.VectorPose())-np.array(self.baseTrans)
            self.base.SetGeoData(basePos,self.chain[0].end.VectorRot())
        self.baseFb.SetGeoData(self.base.VectorPose(),rot=False)
        self.terminalFb.SetGeoData(self.terminal.VectorPose(),rot=False)













